import os
import sys

import fire
import torch
from tqdm import tqdm

project_dir_path = os.path.join(os.path.dirname(__file__), "../../..")
sys.path.append(project_dir_path)

from scripts.datasets.clipping import TextDatasetClippingCondition, TextDatasetClippingHparams, \
    make_text_dataset_clipping_transform
from scripts.datasets.presets import load_preset_llmpt_dataset


def scan_util_rate(
    text_dataset_name: str = "fwe10bt",
    text_dataset_split_name: str = 'train',
    text_dataset_clipping: TextDatasetClippingHparams | None = \
        TextDatasetClippingHparams(
            multiple=False,
            # multiple=True,
            conditions=[
                TextDatasetClippingCondition(text_tokenizer_name="llama3", max_tokens_n=832),
                TextDatasetClippingCondition(text_tokenizer_name="utf8", max_tokens_n=4096)]),
    text_tokenizer_name: str = "llama3",
    chunk_tokens_n: int = 832,
    batch_samples_n: int = 128,
    device: torch.device | str | None = 'cuda',
):
    text_dataset_transform = make_text_dataset_clipping_transform(text_dataset_clipping) \
        if text_dataset_clipping is not None else None
    dataset = load_preset_llmpt_dataset(
        text_dataset_name=text_dataset_name,
        text_dataset_split_name=text_dataset_split_name,
        text_dataset_transform=text_dataset_transform,
        text_tokenizer_name=text_tokenizer_name,
        append_eos=text_dataset_clipping is None,
        batch_samples_n=batch_samples_n,
        chunk_tokens_n=chunk_tokens_n,
        device=device)

    total_pass_tokens_n = 0
    total_real_tokens_n = 0
    progressbar = tqdm(unit='token', unit_scale=True)
    for sample in dataset:
        pass_tokens_n = batch_samples_n * chunk_tokens_n
        real_tokens_n = int((sample.tail - sample.head).sum())
        total_pass_tokens_n += pass_tokens_n
        total_real_tokens_n += real_tokens_n
        progressbar.update(real_tokens_n)
        progressbar.set_postfix({
            'total_pass_tokens_n': total_pass_tokens_n,
            'total_real_tokens_n': total_real_tokens_n,
            'util_rate': total_real_tokens_n / total_pass_tokens_n})


if __name__ == '__main__':
    fire.Fire(scan_util_rate)
